#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
字节跳动 WebSocket TTS 客户端
基于 WebSocket 协议的实时语音合成客户端
"""

import asyncio
import websockets
import uuid
import json
import gzip
import copy
import os
import tempfile
import time
from typing import Optional, Dict, Any
from pathlib import Path
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 消息类型定义
MESSAGE_TYPES = {
    11: "audio-only server response", 
    12: "frontend server response", 
    15: "error message from server"
}

MESSAGE_TYPE_SPECIFIC_FLAGS = {
    0: "no sequence number", 
    1: "sequence number > 0",
    2: "last message from server (seq < 0)", 
    3: "sequence number < 0"
}

MESSAGE_SERIALIZATION_METHODS = {
    0: "no serialization", 
    1: "JSON", 
    15: "custom type"
}

MESSAGE_COMPRESSIONS = {
    0: "no compression", 
    1: "gzip", 
    15: "custom compression method"
}


class TTSWebSocketClient:
    """字节跳动 WebSocket TTS 客户端类（支持连接复用）"""
    
    def __init__(self):
        """
        初始化 WebSocket TTS 客户端
        """
        self.config = self._load_config()
        
        # 默认请求头
        self.default_header = bytearray(b'\x11\x10\x11\x00')
        
        # 连接复用相关
        self._websocket = None
        self._connection_lock = asyncio.Lock()
        self._last_used = 0
        self._connection_timeout = 30  # 30秒后关闭空闲连接
        
        # 取消合成相关
        self._cancel_synthesis = False
        self._current_request_id = None
        
        # 基础请求模板
        self.request_template = {
            "app": {
                "appid": self.config['appid'],
                "token": "access_token",
                "cluster": self.config['cluster']
            },
            "user": {
                "uid": "388808087185088"
            },
            "audio": {
                "voice_type": self.config['voice_type'],
                "encoding": "mp3",
                "speed_ratio": 1.2,  # 稍微加快语速以减少延迟
                "volume_ratio": 1.0,
                "pitch_ratio": 1.0,
            },
            "request": {
                "reqid": "",
                "text": "",
                "text_type": "plain",
                "operation": ""
            }
        }
    
    def _is_connection_open(self, websocket) -> bool:
        """检查WebSocket连接是否打开（兼容不同版本的websockets库）"""
        try:
            # 尝试使用closed属性
            if hasattr(websocket, 'closed'):
                return not websocket.closed
            # 尝试使用state属性
            elif hasattr(websocket, 'state'):
                from websockets.protocol import State
                return websocket.state == State.OPEN
            # 其他情况假设连接已关闭
            else:
                return False
        except:
            return False
    
    def __del__(self):
        """析构函数，确保连接正确关闭"""
        if hasattr(self, '_websocket') and self._websocket and self._is_connection_open(self._websocket):
            try:
                # 在事件循环中安排关闭任务
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    loop.create_task(self._close_connection())
            except:
                pass  # 忽略析构时的异常
    
    def _load_config(self) -> Dict[str, Any]:
        """
        从环境变量加载配置
        
        Returns:
            配置字典
        """
        try:
            config = {
                'appid': os.getenv('TTS_WS_APPID', '2295627400'),
                'token': os.getenv('TTS_WS_TOKEN', 'AnyAoo_B4rPUMpEELgB_N-n2tw7K4OIz'),
                'cluster': os.getenv('TTS_WS_CLUSTER', 'volcano_tts'),
                'voice_type': os.getenv('TTS_WS_VOICE_TYPE', 'BV700_V2_streaming'),
                'host': os.getenv('TTS_WS_HOST', 'openspeech.bytedance.com'),
                'timeout': int(os.getenv('TTS_WS_TIMEOUT', '30')),
                'output_directory': os.getenv('TTS_WS_OUTPUT_DIRECTORY', './output')
            }
            
            # 构建 API URL
            config['api_url'] = os.getenv('TTS_WS_API_URL', f'wss://{config["host"]}/api/v1/tts/ws_binary')
            
            return config
            
        except Exception as e:
            raise ValueError(f"WebSocket TTS 配置加载失败: {e}")
    
    def _build_request(self, text: str, operation: str = "submit") -> bytearray:
        """
        构建 WebSocket 请求数据
        
        Args:
            text: 要合成的文本
            operation: 操作类型 (submit/query)
            
        Returns:
            请求数据
        """
        request_json = copy.deepcopy(self.request_template)
        request_json["request"]["reqid"] = str(uuid.uuid4())
        request_json["request"]["text"] = text
        request_json["request"]["operation"] = operation
        
        # 序列化并压缩
        payload_bytes = str.encode(json.dumps(request_json))
        payload_bytes = gzip.compress(payload_bytes)
        
        # 构建完整请求
        full_request = bytearray(self.default_header)
        full_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
        full_request.extend(payload_bytes)
        
        return full_request
    
    def _parse_response(self, response_data: bytes) -> tuple[bytes, bool]:
        """
        解析 WebSocket 响应数据
        
        Args:
            response_data: 响应数据
            
        Returns:
            (音频数据, 是否完成)
        """
        if len(response_data) < 4:
            return b'', False
            
        # 解析协议头
        protocol_version = response_data[0] >> 4
        header_size = response_data[0] & 0x0f
        message_type = response_data[1] >> 4
        message_type_specific_flags = response_data[1] & 0x0f
        serialization_method = response_data[2] >> 4
        message_compression = response_data[2] & 0x0f
        
        payload = response_data[header_size*4:]
        
        if message_type == 0xb:  # audio-only server response
            if message_type_specific_flags == 0:  # ACK without audio
                return b'', False
            else:
                sequence_number = int.from_bytes(payload[:4], "big", signed=True)
                payload_size = int.from_bytes(payload[4:8], "big", signed=False)
                audio_data = payload[8:]
                
                # 序列号小于0表示最后一个包
                is_done = sequence_number < 0
                return audio_data, is_done
                
        elif message_type == 0xf:  # error message
            code = int.from_bytes(payload[:4], "big", signed=False)
            msg_size = int.from_bytes(payload[4:8], "big", signed=False)
            error_msg = payload[8:]
            if message_compression == 1:
                error_msg = gzip.decompress(error_msg)
            error_msg = str(error_msg, "utf-8")
            raise RuntimeError(f"WebSocket TTS 错误 (代码: {code}): {error_msg}")
            
        return b'', False
    
    async def _get_connection(self):
        """获取或创建WebSocket连接（支持连接复用）"""
        async with self._connection_lock:
            current_time = time.time()
            
            # 检查现有连接是否可用
            if (self._websocket and 
                self._is_connection_open(self._websocket) and 
                current_time - self._last_used < self._connection_timeout):
                self._last_used = current_time
                return self._websocket
            
            # 关闭旧连接
            if self._websocket and self._is_connection_open(self._websocket):
                await self._websocket.close()
            
            # 创建新连接
            headers = {"Authorization": f"Bearer; {self.config['token']}"}
            self._websocket = await websockets.connect(
                self.config['api_url'], 
                additional_headers=headers, 
                ping_interval=None,
                close_timeout=5  # 快速关闭超时
            )
            self._last_used = current_time
            return self._websocket
    
    async def _close_connection(self):
        """关闭WebSocket连接"""
        if self._websocket and self._is_connection_open(self._websocket):
            await self._websocket.close()
            self._websocket = None
    
    async def cancel_synthesis(self):
        """
        取消当前正在进行的语音合成
        """
        self._cancel_synthesis = True
        # 如果有活跃连接，关闭它以立即停止合成
        if self._websocket and self._is_connection_open(self._websocket):
            await self._close_connection()
    
    def reset_cancel_flag(self):
        """
        重置取消标志（在开始新的合成前调用）
        """
        self._cancel_synthesis = False
        self._current_request_id = None
    
    async def synthesize_async(self, text: str, output_file: Optional[str] = None, **kwargs) -> bytes:
        """
        异步合成语音（使用连接复用优化，支持取消）
        
        Args:
            text: 要合成的文本
            output_file: 输出文件路径（可选）
            **kwargs: 其他参数（保持与 TTSClient 兼容）
            
        Returns:
            音频数据（bytes）
            
        Raises:
            asyncio.CancelledError: 当合成被取消时
        """
        if not text.strip():
            raise ValueError("文本内容不能为空")
        
        # 重置取消标志
        self.reset_cancel_flag()
        
        # 生成请求ID用于跟踪
        self._current_request_id = str(uuid.uuid4())
        
        audio_chunks = []
        request_data = self._build_request(text, "submit")
        
        max_retries = 2
        for attempt in range(max_retries):
            try:
                # 检查是否已被取消
                if self._cancel_synthesis:
                    raise asyncio.CancelledError("TTS合成已被取消")
                
                websocket = await self._get_connection()
                
                # 发送请求
                await websocket.send(request_data)
                
                # 接收响应
                while True:
                    # 检查是否已被取消
                    if self._cancel_synthesis:
                        raise asyncio.CancelledError("TTS合成已被取消")
                    
                    try:
                        # 使用超时避免无限等待
                        response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
                    except asyncio.TimeoutError:
                        # 超时时检查取消标志
                        if self._cancel_synthesis:
                            raise asyncio.CancelledError("TTS合成已被取消")
                        continue
                    
                    audio_data, is_done = self._parse_response(response)
                    
                    if audio_data:
                        audio_chunks.append(audio_data)
                    
                    if is_done:
                        break
                
                break  # 成功完成，跳出重试循环
                
            except asyncio.CancelledError:
                # 合成被取消，清理资源
                await self._close_connection()
                raise
            except (websockets.exceptions.ConnectionClosed, 
                    websockets.exceptions.WebSocketException) as e:
                # 连接问题，重置连接并重试
                await self._close_connection()
                if attempt == max_retries - 1:
                    raise RuntimeError(f"WebSocket TTS 合成失败: {e}")
                continue
            except Exception as e:
                raise RuntimeError(f"WebSocket TTS 合成失败: {e}")
        
        # 最后检查是否被取消
        if self._cancel_synthesis:
            raise asyncio.CancelledError("TTS合成已被取消")
        
        # 合并音频数据
        final_audio = b''.join(audio_chunks)
        
        if not final_audio:
            raise ValueError("未获取到有效的音频数据")
        
        # 保存文件（如果指定）
        if output_file:
            self._save_audio_file(final_audio, output_file)
        
        return final_audio
    
    def synthesize(self, text: str, output_file: Optional[str] = None, **kwargs) -> bytes:
        """
        同步合成语音（兼容接口）
        
        Args:
            text: 要合成的文本
            output_file: 输出文件路径（可选）
            **kwargs: 其他参数
            
        Returns:
            音频数据（bytes）
        """
        return asyncio.run(self.synthesize_async(text, output_file, **kwargs))
    
    def _save_audio_file(self, audio_data: bytes, output_file: str) -> None:
        """
        保存音频文件
        
        Args:
            audio_data: 音频数据
            output_file: 输出文件路径
        """
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'wb') as f:
            f.write(audio_data)
    
    def validate_config(self) -> bool:
        """
        验证配置是否有效
        
        Returns:
            配置是否有效
        """
        required_keys = ['appid', 'token', 'cluster', 'voice_type', 'api_url']
        return all(key in self.config and self.config[key] for key in required_keys)
    
    def get_available_speakers(self) -> list:
        """
        获取可用的音色列表
        
        Returns:
            音色列表
        """
        return [
            "BV700_V2_streaming",
            "BV001_streaming",
            "BV002_streaming",
            # 更多音色请参考官方文档
        ]
